package com.light.dp.flink.functions;

import com.light.dp.flink.dto.DataChangeInfo;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.ProcessFunction;
import org.apache.flink.util.Collector;

import java.util.ArrayList;
import java.util.List;

public class BatchProcessFunction extends ProcessFunction<DataChangeInfo, List<DataChangeInfo>> {

    private static final int BATCH_SIZE = 10000;
    private static final long TIME_WINDOW_MS = 5000;
        private transient ValueState<List<DataChangeInfo>> batchState;
        private transient ValueState<Long> startTimeState;

        @Override
        public void open(Configuration parameters) throws Exception {
            ValueStateDescriptor<List<DataChangeInfo>> batchDescriptor = new ValueStateDescriptor<>("batch", new org.apache.flink.api.common.typeinfo.TypeHint<List<DataChangeInfo>>() {}.getTypeInfo());
            batchState = getRuntimeContext().getState(batchDescriptor);

            ValueStateDescriptor<Long> startTimeDescriptor = new ValueStateDescriptor<>("startTime", Long.class);
            startTimeState = getRuntimeContext().getState(startTimeDescriptor);
        }

        @Override
        public void processElement(DataChangeInfo value, Context ctx, Collector<List<DataChangeInfo>> out) throws Exception {
            List<DataChangeInfo> batch = batchState.value();
            if (batch == null) {
                batch = new ArrayList<>();
            }
            batch.add(value);

            Long startTime = startTimeState.value();
            if (startTime == null) {
                startTime = ctx.timerService().currentProcessingTime();
                startTimeState.update(startTime);
            }

            long currentTime = ctx.timerService().currentProcessingTime();
            if (batch.size() >= BATCH_SIZE || currentTime - startTime >= TIME_WINDOW_MS) {
                out.collect(new ArrayList<>(batch));
                batch.clear();
                startTimeState.clear();
            }
            batchState.update(batch);
        }
    }